#!/usr/bin/env python3
"""
YOLO推理服务使用示例
演示如何使用YOLO推理服务进行目标检测
"""

import asyncio
import base64
import json
import time
from pathlib import Path

import cv2
import numpy as np
import requests
from test_client import YOLOClient


def create_sample_image():
    """创建一个示例图像用于测试"""
    # 创建一个640x640的图像
    image = np.zeros((640, 640, 3), dtype=np.uint8)
    
    # 添加一些简单的形状
    # 矩形 (模拟汽车)
    cv2.rectangle(image, (100, 200), (300, 350), (0, 255, 0), -1)
    cv2.putText(image, "CAR", (150, 280), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
    
    # 圆形 (模拟球)
    cv2.circle(image, (500, 150), 50, (255, 0, 0), -1)
    cv2.putText(image, "BALL", (460, 230), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    
    # 椭圆 (模拟人)
    cv2.ellipse(image, (200, 500), (40, 80), 0, 0, 360, (0, 0, 255), -1)
    cv2.putText(image, "PERSON", (150, 450), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    
    return image


def sync_example():
    """同步调用示例"""
    print("=== 同步调用示例 ===")
    
    # 创建示例图像
    image = create_sample_image()
    image_path = "sample_image.jpg"
    cv2.imwrite(image_path, image)
    print(f"创建示例图像: {image_path}")
    
    try:
        # 使用requests进行同步调用
        with open(image_path, 'rb') as f:
            files = {'file': f}
            data = {
                'model_name': 'yolo11n',
                'confidence_threshold': 0.3,
                'iou_threshold': 0.4,
                'draw_results': True
            }
            
            print("发送检测请求...")
            start_time = time.time()
            response = requests.post('http://localhost:8000/detect/file', 
                                   files=files, data=data)
            end_time = time.time()
            
            if response.status_code == 200:
                result = response.json()
                print(f"请求耗时: {(end_time - start_time)*1000:.0f}ms")
                print(f"检测成功: {result['success']}")
                print(f"检测到 {len(result['results'])} 个物体")
                
                # 显示检测结果
                for i, detection in enumerate(result['results']):
                    print(f"  {i+1}. {detection['class_name']}: {detection['confidence']:.2f}")
                
                # 保存结果图像
                if result.get('result_image'):
                    result_image_data = base64.b64decode(result['result_image'])
                    with open('result_sync.jpg', 'wb') as f:
                        f.write(result_image_data)
                    print("结果图像已保存为: result_sync.jpg")
            else:
                print(f"请求失败: {response.status_code} - {response.text}")
                
    except Exception as e:
        print(f"同步调用失败: {e}")
    finally:
        # 清理临时文件
        Path(image_path).unlink(missing_ok=True)


async def async_example():
    """异步调用示例"""
    print("\n=== 异步调用示例 ===")
    
    # 创建示例图像
    image = create_sample_image()
    image_path = "sample_image_async.jpg"
    cv2.imwrite(image_path, image)
    print(f"创建示例图像: {image_path}")
    
    try:
        async with YOLOClient() as client:
            # 检查服务健康状态
            print("检查服务健康状态...")
            health = await client.get_health()
            print(f"服务状态: {health['status']}")
            
            # 列出可用模型
            print("获取模型列表...")
            models = await client.list_models()
            print(f"可用模型数量: {models['total']}")
            
            # 使用文件上传进行检测
            print("执行异步检测...")
            start_time = time.time()
            result = await client.detect_file(
                model_name="yolo11n",
                image_path=image_path,
                confidence_threshold=0.3,
                draw_results=True
            )
            end_time = time.time()
            
            print(f"异步请求耗时: {(end_time - start_time)*1000:.0f}ms")
            print(f"检测成功: {result['success']}")
            print(f"检测到 {len(result['results'])} 个物体")
            
            # 显示检测结果
            for i, detection in enumerate(result['results']):
                print(f"  {i+1}. {detection['class_name']}: {detection['confidence']:.2f}")
                print(f"      位置: ({detection['bbox'][0]:.0f}, {detection['bbox'][1]:.0f}) - "
                      f"({detection['bbox'][2]:.0f}, {detection['bbox'][3]:.0f})")
            
            # 保存结果图像
            if result.get('result_image'):
                result_image = client.decode_base64_to_image(result['result_image'])
                cv2.imwrite('result_async.jpg', result_image)
                print("结果图像已保存为: result_async.jpg")
                
    except Exception as e:
        print(f"异步调用失败: {e}")
    finally:
        # 清理临时文件
        Path(image_path).unlink(missing_ok=True)


async def batch_example():
    """批量检测示例"""
    print("\n=== 批量检测示例 ===")
    
    try:
        # 创建多个示例图像
        images_base64 = []
        for i in range(3):
            # 创建不同的图像
            image = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
            cv2.rectangle(image, (100+i*50, 100+i*50), (200+i*50, 200+i*50), (255, 255, 255), -1)
            
            # 编码为Base64
            _, buffer = cv2.imencode('.jpg', image)
            image_base64 = base64.b64encode(buffer).decode('utf-8')
            images_base64.append(image_base64)
        
        async with YOLOClient() as client:
            print(f"执行批量检测 ({len(images_base64)} 张图像)...")
            start_time = time.time()
            result = await client.batch_detect(
                model_name="yolo11n",
                image_base64_list=images_base64,
                confidence_threshold=0.3
            )
            end_time = time.time()
            
            print(f"批量检测耗时: {(end_time - start_time)*1000:.0f}ms")
            print(f"成功处理: {result['success_count']} 张")
            print(f"处理失败: {result['error_count']} 张")
            
            # 显示每张图像的检测结果
            for item in result['results']:
                index = item['index']
                item_result = item['result']
                if item_result['success']:
                    print(f"  图像 {index+1}: 检测到 {len(item_result['results'])} 个物体")
                else:
                    print(f"  图像 {index+1}: 检测失败 - {item_result['message']}")
                    
    except Exception as e:
        print(f"批量检测失败: {e}")


def base64_example():
    """Base64编码示例"""
    print("\n=== Base64编码示例 ===")
    
    # 创建示例图像
    image = create_sample_image()
    image_path = "sample_base64.jpg"
    cv2.imwrite(image_path, image)
    
    try:
        # 读取图像并编码为Base64
        with open(image_path, 'rb') as f:
            image_data = base64.b64encode(f.read()).decode('utf-8')
        
        # 发送检测请求
        data = {
            "model_name": "yolo11n",
            "image_data": image_data,
            "confidence_threshold": 0.3,
            "iou_threshold": 0.4,
            "draw_results": True,
            "return_crops": False
        }
        
        print("发送Base64检测请求...")
        start_time = time.time()
        response = requests.post('http://localhost:8000/detect', json=data)
        end_time = time.time()
        
        if response.status_code == 200:
            result = response.json()
            print(f"Base64请求耗时: {(end_time - start_time)*1000:.0f}ms")
            print(f"检测成功: {result['success']}")
            print(f"检测到 {len(result['results'])} 个物体")
            
            # 保存结果图像
            if result.get('result_image'):
                result_image_data = base64.b64decode(result['result_image'])
                with open('result_base64.jpg', 'wb') as f:
                    f.write(result_image_data)
                print("结果图像已保存为: result_base64.jpg")
        else:
            print(f"Base64请求失败: {response.status_code} - {response.text}")
            
    except Exception as e:
        print(f"Base64示例失败: {e}")
    finally:
        # 清理临时文件
        Path(image_path).unlink(missing_ok=True)


def main():
    """主函数"""
    print("YOLO推理服务使用示例")
    print("=" * 50)
    print("请确保YOLO推理服务已启动 (http://localhost:8000)")
    print()
    
    # 检查服务是否可用
    try:
        response = requests.get('http://localhost:8000/health', timeout=5)
        if response.status_code == 200:
            print("✅ 服务连接正常")
        else:
            print("❌ 服务连接失败")
            return
    except Exception as e:
        print(f"❌ 无法连接到服务: {e}")
        print("请先启动YOLO推理服务:")
        print("  python yolo_inference_server.py")
        print("  或")
        print("  ./start_server.sh")
        return
    
    # 运行示例
    sync_example()
    asyncio.run(async_example())
    asyncio.run(batch_example())
    base64_example()
    
    print("\n✅ 所有示例运行完成！")
    print("生成的结果图像:")
    print("  - result_sync.jpg (同步调用结果)")
    print("  - result_async.jpg (异步调用结果)")
    print("  - result_base64.jpg (Base64调用结果)")


if __name__ == "__main__":
    main()
